package com.ajocer.springbootinit.rag;

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.util.Arrays;
import java.util.HashSet;
import java.util.Set;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * 提示词压缩工具类
 * 用于优化提示词，减少Token消耗
 */
@Component
@Slf4j
public class PromptCompressor {

    // 冗余短语模式
    private static final Pattern REDUNDANT_PATTERN = Pattern.compile(
            "(?i)(please|kindly|would you|could you|I would like you to|I want you to|make sure to|remember to)\\s");
    
    // 重复指令模式
    private static final Pattern DUPLICATE_INSTRUCTION_PATTERN = Pattern.compile(
            "(\\b(?:provide|give|list|show|tell|explain|describe|summarize|analyze)\\b[^.!?]{10,100})[.!?]\\s+\\1", 
            Pattern.CASE_INSENSITIVE);
    
    // 多余空白字符
    private static final Pattern WHITESPACE_PATTERN = Pattern.compile("\\s+");
    
    // 常见的停用词
    private static final Set<String> STOP_WORDS = new HashSet<>(Arrays.asList(
            "a", "an", "the", "and", "or", "but", "if", "then", "else", "when",
            "up", "down", "in", "out", "on", "off", "over", "under", "again",
            "further", "then", "once", "here", "there", "why", "how", "all", 
            "any", "both", "each", "few", "more", "most", "other", "some", "such",
            "no", "nor", "not", "only", "own", "same", "so", "than", "too", "very"
    ));

    /**
     * 压缩提示词，减少Token消耗
     * @param prompt 原始提示词
     * @return 压缩后的提示词
     */
    public String compress(String prompt) {
        if (prompt == null || prompt.isEmpty()) {
            return prompt;
        }
        
        try {
            // 记录原始Token估算
            int originalTokens = estimateTokens(prompt);
            
            // 1. 删除冗余短语
            String compressed = removeRedundantPhrases(prompt);
            
            // 2. 删除重复指令
            compressed = removeDuplicateInstructions(compressed);
            
            // 3. 规范化空白字符
            compressed = normalizeWhitespace(compressed);
            
            // 4. 简化示例（如果有）
            compressed = simplifyExamples(compressed);
            
            // 记录压缩后Token估算
            int compressedTokens = estimateTokens(compressed);
            
            // 计算节省比例
            double savingRatio = (originalTokens > 0) 
                    ? (double)(originalTokens - compressedTokens) / originalTokens * 100 
                    : 0;
            
            log.debug("提示词压缩: {} -> {} tokens (节省 {:.1f}%)", 
                    originalTokens, compressedTokens, savingRatio);
            
            return compressed;
        } catch (Exception e) {
            // 如果压缩过程出错，返回原始提示词
            log.warn("提示词压缩失败，使用原始提示词", e);
            return prompt;
        }
    }
    
    /**
     * 删除冗余短语
     */
    private String removeRedundantPhrases(String prompt) {
        return REDUNDANT_PATTERN.matcher(prompt).replaceAll("");
    }
    
    /**
     * 删除重复指令
     */
    private String removeDuplicateInstructions(String prompt) {
        Matcher matcher = DUPLICATE_INSTRUCTION_PATTERN.matcher(prompt);
        return matcher.replaceAll("$1");
    }
    
    /**
     * 规范化空白字符
     */
    private String normalizeWhitespace(String prompt) {
        return WHITESPACE_PATTERN.matcher(prompt).replaceAll(" ").trim();
    }
    
    /**
     * 简化示例
     */
    private String simplifyExamples(String prompt) {
        // 识别示例区块
        Pattern examplePattern = Pattern.compile("(?s)(?:Example|For example|Here's an example)[:\\s]+(.*?)(?=\\n\\n|$)");
        Matcher matcher = examplePattern.matcher(prompt);
        
        StringBuffer result = new StringBuffer();
        while (matcher.find()) {
            String example = matcher.group(1);
            // 如果示例太长，进行简化
            if (example.length() > 200) {
                String simplified = example.substring(0, 150) + "... [example simplified]";
                matcher.appendReplacement(result, "Example: " + Matcher.quoteReplacement(simplified));
            } else {
                matcher.appendReplacement(result, "Example: " + Matcher.quoteReplacement(example));
            }
        }
        matcher.appendTail(result);
        
        return result.toString();
    }
    
    /**
     * 估算Token数量（粗略估计）
     * 基于GPT模型的分词规则，大约4个字符为1个token
     */
    private int estimateTokens(String text) {
        if (text == null || text.isEmpty()) {
            return 0;
        }
        // 粗略估计：每4个字符约1个token
        return text.length() / 4;
    }
} 